{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9e9fe848",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/tensorrt_torchtrt_efficientnet/nvidia_logo.png\" width=\"90px\">\n",
    "\n",
    "# PySpark Huggingface Inferencing\n",
    "### Sentiment Analysis using Pipelines with Tensorflow\n",
    "\n",
    "In this notebook, we demonstrate distributed inference with Huggingface Pipelines to perform sentiment analysis.  \n",
    "From: https://huggingface.co/docs/transformers/quicktour#pipeline-usage"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1799fd4f",
   "metadata": {},
   "source": [
    "Note that cuFFT/cuDNN/cuBLAS registration errors are expected (as of `tf=2.17.0`) and will not affect behavior, as noted in [this issue.](https://github.com/tensorflow/tensorflow/issues/62075)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0dd0f77b-ee1b-4477-a038-d25a4f1da0ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 13:57:08.242673: I tensorflow/core/util/port.cc:153] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-02-04 13:57:08.249833: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2025-02-04 13:57:08.257735: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2025-02-04 13:57:08.259994: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2025-02-04 13:57:08.266655: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2025-02-04 13:57:08.649929: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "from transformers import pipeline\n",
    "\n",
    "# Manually enable Huggingface tokenizer parallelism to avoid disabling with PySpark parallelism.\n",
    "# See (https://github.com/huggingface/transformers/issues/5486) for more info. \n",
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d80fc3f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "I0000 00:00:1738706229.309141 3668091 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706229.333555 3668091 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706229.336487 3668091 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n"
     ]
    }
   ],
   "source": [
    "device = 0 if tf.config.list_physical_devices('GPU') else -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e60a2877",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.17.0\n"
     ]
    }
   ],
   "source": [
    "# Enable GPU memory growth\n",
    "gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "if gpus:\n",
    "    try:\n",
    "        for gpu in gpus:\n",
    "            tf.config.experimental.set_memory_growth(gpu, True)\n",
    "    except RuntimeError as e:\n",
    "        print(e)\n",
    "\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "553b28d2-a5d1-4d07-8a49-8f82b808e738",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No model was supplied, defaulted to distilbert/distilbert-base-uncased-finetuned-sst-2-english and revision 714eb0f (https://huggingface.co/distilbert/distilbert-base-uncased-finetuned-sst-2-english).\n",
      "Using a pipeline without specifying a model name and revision in production is not recommended.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "I0000 00:00:1738706229.617170 3668091 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706229.620218 3668091 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706229.622781 3668091 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706229.732012 3668091 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706229.733045 3668091 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "I0000 00:00:1738706229.733965 3668091 cuda_executor.cc:1015] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355\n",
      "2025-02-04 13:57:09.734873: I tensorflow/core/common_runtime/gpu/gpu_device.cc:2021] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 43096 MB memory:  -> device: 0, name: NVIDIA RTX A6000, pci bus id: 0000:01:00.0, compute capability: 8.6\n",
      "All PyTorch model weights were used when initializing TFDistilBertForSequenceClassification.\n",
      "\n",
      "All the weights of TFDistilBertForSequenceClassification were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFDistilBertForSequenceClassification for predictions without further training.\n"
     ]
    }
   ],
   "source": [
    "classifier = pipeline(\"sentiment-analysis\", device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b91fe91-b725-4564-ae93-56e3fb51e47c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'label': 'POSITIVE', 'score': 0.9997794032096863}]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classifier((\"We are very happy to show you the 🤗 Transformers library.\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0be39eb3-462c-42ff-b8f4-09f4e4fe3a3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "label: POSITIVE, with score: 0.9998\n",
      "label: NEGATIVE, with score: 0.5282\n"
     ]
    }
   ],
   "source": [
    "results = classifier([\"We are very happy to show you the 🤗 Transformers library.\", \"We hope you don't hate it.\"])\n",
    "for result in results:\n",
    "    print(f\"label: {result['label']}, with score: {round(result['score'], 4)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e29ee6d8",
   "metadata": {},
   "source": [
    "Let's try a different model and tokenizer in the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "cd9d3349",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"nlptown/bert-base-multilingual-uncased-sentiment\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "99e21b58",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "All PyTorch model weights were used when initializing TFBertForSequenceClassification.\n",
      "\n",
      "All the weights of TFBertForSequenceClassification were initialized from the PyTorch model.\n",
      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFBertForSequenceClassification for predictions without further training.\n"
     ]
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, TFAutoModelForSequenceClassification\n",
    "\n",
    "model = TFAutoModelForSequenceClassification.from_pretrained(model_name)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31079133",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'label': '5 stars', 'score': 0.7272477746009827}]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "classifier = pipeline(\"sentiment-analysis\", model=model, tokenizer=tokenizer, device=device)\n",
    "classifier(\"Nous sommes très heureux de vous présenter la bibliothèque 🤗 Transformers.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6357234",
   "metadata": {},
   "source": [
    "## PySpark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "69dd6a1a-f450-47f0-9dbf-ad250585a011",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, struct, pandas_udf\n",
    "from pyspark.ml.functions import predict_batch_udf\n",
    "from pyspark.sql.types import *\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark import SparkConf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "287b1e96",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import pandas as pd\n",
    "import datasets\n",
    "from datasets import load_dataset\n",
    "datasets.disable_progress_bars()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50e124cd",
   "metadata": {},
   "source": [
    "Check the cluster environment to handle any platform-specific Spark configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "36001f55",
   "metadata": {},
   "outputs": [],
   "source": [
    "on_databricks = os.environ.get(\"DATABRICKS_RUNTIME_VERSION\", False)\n",
    "on_dataproc = os.environ.get(\"DATAPROC_IMAGE_VERSION\", False)\n",
    "on_standalone = not (on_databricks or on_dataproc)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48c7271a",
   "metadata": {},
   "source": [
    "#### Create Spark Session\n",
    "\n",
    "For local standalone clusters, we'll connect to the cluster and create the Spark Session.  \n",
    "For CSP environments, Spark will either be preconfigured (Databricks) or we'll need to create the Spark Session (Dataproc)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6e0e0dd7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:57:12 WARN Utils: Your hostname, cb4ae00-lcedt resolves to a loopback address: 127.0.1.1; using 10.110.47.100 instead (on interface eno1)\n",
      "25/02/04 13:57:12 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address\n",
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
      "25/02/04 13:57:12 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n"
     ]
    }
   ],
   "source": [
    "conf = SparkConf()\n",
    "\n",
    "if 'spark' not in globals():\n",
    "    if on_standalone:\n",
    "        import socket\n",
    "        \n",
    "        conda_env = os.environ.get(\"CONDA_PREFIX\")\n",
    "        hostname = socket.gethostname()\n",
    "        conf.setMaster(f\"spark://{hostname}:7077\")\n",
    "        conf.set(\"spark.pyspark.python\", f\"{conda_env}/bin/python\")\n",
    "        conf.set(\"spark.pyspark.driver.python\", f\"{conda_env}/bin/python\")\n",
    "    elif on_dataproc:\n",
    "        conf.set(\"spark.executorEnv.TF_GPU_ALLOCATOR\", \"cuda_malloc_async\")\n",
    "\n",
    "    conf.set(\"spark.executor.cores\", \"8\")\n",
    "    conf.set(\"spark.task.resource.gpu.amount\", \"0.125\")\n",
    "    conf.set(\"spark.executor.resource.gpu.amount\", \"1\")\n",
    "    conf.set(\"spark.sql.execution.arrow.pyspark.enabled\", \"true\")\n",
    "    conf.set(\"spark.python.worker.reuse\", \"true\")\n",
    "\n",
    "conf.set(\"spark.sql.execution.arrow.maxRecordsPerBatch\", \"1000\")\n",
    "spark = SparkSession.builder.appName(\"spark-dl-examples\").config(conf=conf).getOrCreate()\n",
    "sc = spark.sparkContext"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "42d70208",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"imdb\", split=\"test\")\n",
    "dataset = dataset.to_pandas().drop(columns=\"label\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95ded4b2",
   "metadata": {},
   "source": [
    "#### Create PySpark DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ac24f3c2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StructType([StructField('text', StringType(), True)])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = spark.createDataFrame(dataset).repartition(8)\n",
    "df.schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "1db4db3a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "25000"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "517fe2e9",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:57:20 WARN TaskSetManager: Stage 6 contains a task of very large size (4021 KiB). The maximum recommended task size is 1000 KiB.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[Row(text=\"Anyone remember the first CKY, CKY2K etc..? Back when it was about making crazy cool stuff, rather than watching Bam Margera act like a douchebag, spoiled 5 year old, super/rock-star wannabe.<br /><br />The show used to be awesome, however, Bam's fame and wealth has led him to believe, that we now enjoy him acting childish and idiotic, more than actual cool stuff, that used to be in ex. CKY2K.<br /><br />The acts are so repetitive, there's like nothing new, except annoying stupidity and rehearsed comments... The only things we see is Bam Margera, so busy showing us how much he doesn't care, how much money he got or whatsoever.<br /><br />I really got nothing much left to say except, give us back CKY2K, cause Bam suck..<br /><br />I enjoy watching Steve-o, Knoxville etc. a thousand times more.\")]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.take(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "e176d28b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:57:20 WARN TaskSetManager: Stage 9 contains a task of very large size (4021 KiB). The maximum recommended task size is 1000 KiB.\n"
     ]
    }
   ],
   "source": [
    "data_path = \"spark-dl-datasets/imdb_test\"\n",
    "if on_databricks:\n",
    "    dbutils.fs.mkdirs(\"/FileStore/spark-dl-datasets\")\n",
    "    data_path = \"dbfs:/FileStore/\" + data_path\n",
    "\n",
    "df.write.mode(\"overwrite\").parquet(data_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "395e0374",
   "metadata": {},
   "source": [
    "#### Load and preprocess DataFrame\n",
    "\n",
    "Define our preprocess function. We'll take the first sentence from each sample as our input for sentiment analysis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "9665b7b6-d7e9-4bd4-b29d-7a449ac5b574",
   "metadata": {},
   "outputs": [],
   "source": [
    "@pandas_udf(\"string\")\n",
    "def preprocess(text: pd.Series) -> pd.Series:\n",
    "    return pd.Series([s.split(\".\")[0] for s in text])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "26693020",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------------------------------------------------------------------------------------------------+\n",
      "|                                                                                               input|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "|Doesn't anyone bother to check where this kind of sludge comes from before blathering on about it...|\n",
      "|                          There were two things I hated about WASTED : The directing and the script |\n",
      "|                                I'm rather surprised that anybody found this film touching or moving|\n",
      "|Cultural Vandalism Is the new Hallmark production of Gulliver's Travels an act of cultural vandal...|\n",
      "|I was at Wrestlemania VI in Toronto as a 10 year old, and the event I saw then was pretty differe...|\n",
      "|                                                                     This movie has been done before|\n",
      "|[ as a new resolution for this year 2005, i decide to write a comment for each movie I saw in the...|\n",
      "|This movie is over hyped!! I am sad to say that I manage to watch the first 15 minutes of this mo...|\n",
      "|This show had a promising start as sort of the opposite of 'Oceans 11' but has developed into a s...|\n",
      "|MINOR PLOT SPOILERS AHEAD!!!<br /><br />How did such talented actors get involved in such mindles...|\n",
      "|                              There is not one character on this sitcom with any redeeming qualities|\n",
      "|                                  Tommy Lee Jones was the best Woodroe and no one can play Woodroe F|\n",
      "|                                 My wife rented this movie and then conveniently never got to see it|\n",
      "|This is one of those star-filled over-the-top comedies that could a) be hysterical, or b) wish th...|\n",
      "|This excruciatingly boring and unfunny movie made me think that Chaplin was the real Hitler, as o...|\n",
      "|                           you will likely be sorely disappointed by this sequel that's not a sequel|\n",
      "|                          If I was British, I would be embarrassed by this portrayal of incompetence|\n",
      "|One of those movies in which there are no big twists whatsoever and you can predict pretty much w...|\n",
      "|                        This show is like watching someone who is in training to someday host a show|\n",
      "|                                                                                                Sigh|\n",
      "+----------------------------------------------------------------------------------------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Limit to N rows, since this can be slow\n",
    "df = spark.read.parquet(data_path).limit(256).repartition(8)\n",
    "df = df.select(preprocess(col(\"text\")).alias(\"input\")).cache()\n",
    "df.show(truncate=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76dc525c",
   "metadata": {},
   "source": [
    "## Inference using Spark DL API\n",
    "\n",
    "Distributed inference using the PySpark [predict_batch_udf](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.functions.predict_batch_udf.html#pyspark.ml.functions.predict_batch_udf):\n",
    "\n",
    "- predict_batch_fn uses Tensorflow APIs to load the model and return a predict function which operates on numpy arrays \n",
    "- predict_batch_udf will convert the Spark DataFrame columns into numpy input batches for the predict function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "0da9d25c-5ebe-4503-bb19-154fcc047cbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_batch_fn():\n",
    "    import tensorflow as tf\n",
    "    from transformers import pipeline\n",
    "\n",
    "    # Enable GPU memory growth\n",
    "    gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        try:\n",
    "            for gpu in gpus:\n",
    "                tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        except RuntimeError as e:\n",
    "            print(e)\n",
    "    \n",
    "    device = 0 if tf.config.list_physical_devices('GPU') else -1\n",
    "    pipe = pipeline(\"sentiment-analysis\", device=device)\n",
    "    def predict(inputs):\n",
    "        return pipe(inputs.tolist())\n",
    "    return predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "78afef29-ee30-4267-9fb6-be2dcb86cbba",
   "metadata": {},
   "outputs": [],
   "source": [
    "classify = predict_batch_udf(predict_batch_fn,\n",
    "                             return_type=StructType([\n",
    "                                 StructField(\"label\", StringType(), True),\n",
    "                                 StructField(\"score\", FloatType(), True)\n",
    "                             ]),\n",
    "                             batch_size=32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "a5bc327e-89cf-4731-82e6-e66cb93deef1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 18:=======>                                                  (1 + 7) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 8.06 ms, sys: 2.92 ms, total: 11 ms\n",
      "Wall time: 4.86 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "# note: expanding the \"struct\" return_type to top-level columns\n",
    "preds = df.withColumn(\"preds\", classify(struct(\"input\"))).select(\"input\", \"preds.*\")\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "ac642895-cfd6-47ee-9b21-02e7835424e4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 21:>                                                         (0 + 8) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.5 ms, sys: 1.43 ms, total: 5.93 ms\n",
      "Wall time: 1.19 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", classify(\"input\")).select(\"input\", \"preds.*\")\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "76a44d80-d5db-405f-989c-7246379cfb95",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 24:>                                                         (0 + 8) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.9 ms, sys: 605 μs, total: 6.5 ms\n",
      "Wall time: 1.37 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", classify(col(\"input\"))).select(\"input\", \"preds.*\")\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "c01761b3-c766-46b0-ae0b-fcf968ffb3a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 27:>                                                         (0 + 1) / 1]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------------------------------------------------------------+--------+----------+\n",
      "|                                                                           input|   label|     score|\n",
      "+--------------------------------------------------------------------------------+--------+----------+\n",
      "|Doesn't anyone bother to check where this kind of sludge comes from before bl...|NEGATIVE| 0.9984061|\n",
      "|      There were two things I hated about WASTED : The directing and the script |NEGATIVE| 0.9979007|\n",
      "|            I'm rather surprised that anybody found this film touching or moving|POSITIVE|0.83874947|\n",
      "|Cultural Vandalism Is the new Hallmark production of Gulliver's Travels an ac...|NEGATIVE|0.99727434|\n",
      "|I was at Wrestlemania VI in Toronto as a 10 year old, and the event I saw the...|POSITIVE|  0.982114|\n",
      "|                                                 This movie has been done before|NEGATIVE|0.94210696|\n",
      "|[ as a new resolution for this year 2005, i decide to write a comment for eac...|NEGATIVE| 0.9967818|\n",
      "|This movie is over hyped!! I am sad to say that I manage to watch the first 1...|NEGATIVE| 0.9985843|\n",
      "|This show had a promising start as sort of the opposite of 'Oceans 11' but ha...|NEGATIVE|0.99926835|\n",
      "|MINOR PLOT SPOILERS AHEAD!!!<br /><br />How did such talented actors get invo...|NEGATIVE|0.99956733|\n",
      "|          There is not one character on this sitcom with any redeeming qualities|NEGATIVE| 0.9985662|\n",
      "|              Tommy Lee Jones was the best Woodroe and no one can play Woodroe F|POSITIVE|  0.994562|\n",
      "|             My wife rented this movie and then conveniently never got to see it|NEGATIVE|0.99841607|\n",
      "|This is one of those star-filled over-the-top comedies that could a) be hyste...|NEGATIVE| 0.9953243|\n",
      "|This excruciatingly boring and unfunny movie made me think that Chaplin was t...|NEGATIVE| 0.9997607|\n",
      "|       you will likely be sorely disappointed by this sequel that's not a sequel|NEGATIVE| 0.9997198|\n",
      "|      If I was British, I would be embarrassed by this portrayal of incompetence|NEGATIVE| 0.9965172|\n",
      "|One of those movies in which there are no big twists whatsoever and you can p...|NEGATIVE| 0.9986059|\n",
      "|    This show is like watching someone who is in training to someday host a show|NEGATIVE|0.97015846|\n",
      "|                                                                            Sigh|NEGATIVE| 0.9923151|\n",
      "+--------------------------------------------------------------------------------+--------+----------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "preds.show(truncate=80)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc8127d9",
   "metadata": {},
   "source": [
    "## Using Triton Inference Server\n",
    "In this section, we demonstrate integration with the [Triton Inference Server](https://developer.nvidia.com/nvidia-triton-inference-server), an open-source, GPU-accelerated serving solution for DL.  \n",
    "We use [PyTriton](https://github.com/triton-inference-server/pytriton), a Flask-like framework that handles client/server communication with the Triton server.  \n",
    "\n",
    "The process looks like this:\n",
    "- Distribute a PyTriton task across the Spark cluster, instructing each node to launch a Triton server process.\n",
    "- Define a Triton inference function, which contains a client that binds to the local server on a given node and sends inference requests.\n",
    "- Wrap the Triton inference function in a predict_batch_udf to launch parallel inference requests using Spark.\n",
    "- Finally, distribute a shutdown signal to terminate the Triton server processes on each node.\n",
    "\n",
    "<img src=\"../images/spark-server.png\" alt=\"drawing\" width=\"700\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "4d4be844-4b8c-47df-bd09-0c280c7ff16b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f15dfcb",
   "metadata": {},
   "source": [
    "Import the helper class from server_utils.py:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "bfa7ec9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "sc.addPyFile(\"server_utils.py\")\n",
    "\n",
    "from server_utils import TritonServerManager"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bf04546",
   "metadata": {},
   "source": [
    "Define the Triton Server function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "7e53df9f-43cb-4c38-b8ac-dc2cbad99815",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triton_server(ports):\n",
    "    import time\n",
    "    import signal\n",
    "    import numpy as np\n",
    "    import tensorflow as tf\n",
    "    from transformers import pipeline\n",
    "    from pytriton.decorators import batch\n",
    "    from pytriton.model_config import DynamicBatcher, ModelConfig, Tensor\n",
    "    from pytriton.triton import Triton, TritonConfig\n",
    "    from pyspark import TaskContext\n",
    "\n",
    "    print(f\"SERVER: Initializing pipeline on worker {TaskContext.get().partitionId()}.\")\n",
    "    # Enable GPU memory growth\n",
    "    gpus = tf.config.experimental.list_physical_devices('GPU')\n",
    "    if gpus:\n",
    "        try:\n",
    "            for gpu in gpus:\n",
    "                tf.config.experimental.set_memory_growth(gpu, True)\n",
    "        except RuntimeError as e:\n",
    "            print(e)\n",
    "    \n",
    "    device = 0 if tf.config.list_physical_devices('GPU') else -1\n",
    "    \n",
    "    pipe = pipeline(\"sentiment-analysis\", device=device)\n",
    "    print(f\"SERVER: Using {device} device.\")\n",
    "\n",
    "    @batch\n",
    "    def _infer_fn(**inputs):\n",
    "        sentences = np.squeeze(inputs[\"text\"]).tolist()\n",
    "        print(f\"SERVER: Received batch of size {len(sentences)}\")\n",
    "        decoded_sentences = [s.decode(\"utf-8\") for s in sentences]\n",
    "        return {\n",
    "            \"outputs\": np.array([[json.dumps(o)] for o in pipe(decoded_sentences)])\n",
    "        }\n",
    "\n",
    "    workspace_path = f\"/tmp/triton_{time.strftime('%m_%d_%M_%S')}\"\n",
    "    triton_conf = TritonConfig(http_port=ports[0], grpc_port=ports[1], metrics_port=ports[2])\n",
    "    with Triton(config=triton_conf, workspace=workspace_path) as triton:\n",
    "        triton.bind(\n",
    "            model_name=\"SentimentAnalysis\",\n",
    "            infer_func=_infer_fn,\n",
    "            inputs=[\n",
    "                Tensor(name=\"text\", dtype=object, shape=(-1,)),\n",
    "            ],\n",
    "            outputs=[\n",
    "                Tensor(name=\"outputs\", dtype=object, shape=(-1,)),\n",
    "            ],\n",
    "            config=ModelConfig(\n",
    "                max_batch_size=64,\n",
    "                batcher=DynamicBatcher(max_queue_delay_microseconds=5000),  # 5ms\n",
    "            ),\n",
    "            strict=True,\n",
    "        )\n",
    "\n",
    "        def _stop_triton(signum, frame):\n",
    "            print(\"SERVER: Received SIGTERM. Stopping Triton server.\")\n",
    "            triton.stop()\n",
    "\n",
    "        signal.signal(signal.SIGTERM, _stop_triton)\n",
    "\n",
    "        print(\"SERVER: Serving inference\")\n",
    "        triton.serve()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19d9028d",
   "metadata": {},
   "source": [
    "#### Start Triton servers"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5354c597",
   "metadata": {},
   "source": [
    "The `TritonServerManager` will handle the lifecycle of Triton server instances across the Spark cluster:\n",
    "- Find available ports for HTTP/gRPC/metrics\n",
    "- Deploy a server on each node via stage-level scheduling\n",
    "- Gracefully shutdown servers across nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "156de815",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"SentimentAnalysis\"\n",
    "server_manager = TritonServerManager(model_name=model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d003a862",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-07 11:03:44,809 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-07 11:03:44,810 - INFO - Starting 1 servers.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'cb4ae00-lcedt': (2020631, [7000, 7001, 7002])}"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Returns {'hostname', (server_pid, [http_port, grpc_port, metrics_port])}\n",
    "server_manager.start_servers(triton_server)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4c4017c",
   "metadata": {},
   "source": [
    "#### Define client function"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "405edc49",
   "metadata": {},
   "source": [
    "Get the hostname -> url mapping from the server manager:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19768ddb",
   "metadata": {},
   "outputs": [],
   "source": [
    "host_to_http_url = server_manager.host_to_http_url  # or server_manager.host_to_grpc_url"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb5dbb89",
   "metadata": {},
   "source": [
    "Define the Triton inference function, which returns a predict function for batch inference through the server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "431b864c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def triton_fn(model_name, host_to_url):\n",
    "    import socket\n",
    "    import numpy as np\n",
    "    from pytriton.client import ModelClient\n",
    "\n",
    "    url = host_to_url[socket.gethostname()]\n",
    "    print(f\"Connecting to Triton model {model_name} at {url}.\")\n",
    "\n",
    "    def infer_batch(inputs):\n",
    "        with ModelClient(url, model_name, inference_timeout_s=240) as client:\n",
    "            flattened = np.squeeze(inputs).tolist()\n",
    "            # Encode batch\n",
    "            encoded_batch = [[text.encode(\"utf-8\")] for text in flattened]\n",
    "            encoded_batch_np = np.array(encoded_batch, dtype=np.bytes_)\n",
    "            # Run inference\n",
    "            result_data = client.infer_batch(encoded_batch_np)\n",
    "            result_data = np.squeeze(result_data[\"outputs\"], -1)\n",
    "            return [json.loads(o) for o in result_data]\n",
    "        \n",
    "    return infer_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "3930cfcd-3284-4c6a-a9b5-36b8053fe899",
   "metadata": {},
   "outputs": [],
   "source": [
    "classify = predict_batch_udf(partial(triton_fn, model_name=model_name, host_to_url=host_to_http_url),\n",
    "                             return_type=StructType([\n",
    "                                 StructField(\"label\", StringType(), True),\n",
    "                                 StructField(\"score\", FloatType(), True)\n",
    "                             ]),\n",
    "                             input_tensor_shapes=[[1]],\n",
    "                             batch_size=32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a8ec7be",
   "metadata": {},
   "source": [
    "#### Load and preprocess DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "d53fb283-bf9e-4571-8c68-b75a41f1f067",
   "metadata": {},
   "outputs": [],
   "source": [
    "@pandas_udf(\"string\")\n",
    "def preprocess(text: pd.Series) -> pd.Series:\n",
    "    return pd.Series([s.split(\".\")[0] for s in text])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "29b0cc0d-c480-4e4a-bd41-207dc314cba5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "25/02/04 13:57:36 WARN CacheManager: Asked to cache already cached data.\n"
     ]
    }
   ],
   "source": [
    "df = spark.read.parquet(data_path).limit(256).repartition(8)\n",
    "df = df.select(preprocess(col(\"text\")).alias(\"input\")).cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da39990f",
   "metadata": {},
   "source": [
    "#### Run Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "8eecbf23-4e9e-4d4c-8645-98209b25db2c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 33:===========================================>              (6 + 2) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 13.1 ms, sys: 8.29 ms, total: 21.4 ms\n",
      "Wall time: 7.54 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "# first pass caches model/fn\n",
    "# note: expanding the \"struct\" return_type to top-level columns\n",
    "preds = df.withColumn(\"preds\", classify(struct(\"input\"))).select(\"input\", \"preds.*\")\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "566ba28c-0ca4-4479-a24a-c8a362228b89",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 36:===========================================>              (6 + 2) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.54 ms, sys: 3.13 ms, total: 10.7 ms\n",
      "Wall time: 7.02 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", classify(\"input\")).select(\"input\", \"preds.*\")\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "44c7e776-08da-484a-ba07-9d6add1a0f15",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 39:===========================================>              (6 + 2) / 8]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 6.26 ms, sys: 3 ms, total: 9.26 ms\n",
      "Wall time: 7.03 s\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "preds = df.withColumn(\"preds\", classify(col(\"input\"))).select(\"input\", \"preds.*\")\n",
    "results = preds.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "f61d79f8-661e-4d9e-a3aa-c0754b854603",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------------------------------------------------------------------------+--------+----------+\n",
      "|                                                                           input|   label|     score|\n",
      "+--------------------------------------------------------------------------------+--------+----------+\n",
      "|Doesn't anyone bother to check where this kind of sludge comes from before bl...|NEGATIVE| 0.9984061|\n",
      "|      There were two things I hated about WASTED : The directing and the script |NEGATIVE| 0.9979007|\n",
      "|            I'm rather surprised that anybody found this film touching or moving|POSITIVE|0.83874947|\n",
      "|Cultural Vandalism Is the new Hallmark production of Gulliver's Travels an ac...|NEGATIVE|0.99727434|\n",
      "|I was at Wrestlemania VI in Toronto as a 10 year old, and the event I saw the...|POSITIVE|  0.982114|\n",
      "|                                                 This movie has been done before|NEGATIVE|0.94210696|\n",
      "|[ as a new resolution for this year 2005, i decide to write a comment for eac...|NEGATIVE| 0.9967818|\n",
      "|This movie is over hyped!! I am sad to say that I manage to watch the first 1...|NEGATIVE| 0.9985843|\n",
      "|This show had a promising start as sort of the opposite of 'Oceans 11' but ha...|NEGATIVE|0.99926835|\n",
      "|MINOR PLOT SPOILERS AHEAD!!!<br /><br />How did such talented actors get invo...|NEGATIVE|0.99956733|\n",
      "|          There is not one character on this sitcom with any redeeming qualities|NEGATIVE| 0.9985662|\n",
      "|              Tommy Lee Jones was the best Woodroe and no one can play Woodroe F|POSITIVE|  0.994562|\n",
      "|             My wife rented this movie and then conveniently never got to see it|NEGATIVE|0.99841607|\n",
      "|This is one of those star-filled over-the-top comedies that could a) be hyste...|NEGATIVE| 0.9953243|\n",
      "|This excruciatingly boring and unfunny movie made me think that Chaplin was t...|NEGATIVE| 0.9997607|\n",
      "|       you will likely be sorely disappointed by this sequel that's not a sequel|NEGATIVE| 0.9997198|\n",
      "|      If I was British, I would be embarrassed by this portrayal of incompetence|NEGATIVE| 0.9965172|\n",
      "|One of those movies in which there are no big twists whatsoever and you can p...|NEGATIVE| 0.9986059|\n",
      "|    This show is like watching someone who is in training to someday host a show|NEGATIVE|0.97015846|\n",
      "|                                                                            Sigh|NEGATIVE| 0.9923151|\n",
      "+--------------------------------------------------------------------------------+--------+----------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "preds.show(truncate=80)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fac2ae57",
   "metadata": {},
   "source": [
    "#### Shut down server on each executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "425d3b28-7705-45ba-8a18-ad34fc895219",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-02-04 13:57:58,747 - INFO - Requesting stage-level resources: (cores=5, gpu=1.0)\n",
      "2025-02-04 13:58:03,931 - INFO - Sucessfully stopped 1 servers.                 \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[True]"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "server_manager.stop_servers()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "9f19643c-4ee4-44f2-b762-2078c0c8eba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not on_databricks: # on databricks, spark.stop() puts the cluster in a bad state\n",
    "    spark.stop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a538c47-317d-4cac-b9b9-559e88677518",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "spark-dl-tf",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
